#include "Process.h"
#include <psapi.h>

BOOL CProcess::EnableDebugPrivilege()
{
	HANDLE hToken;
	LUID sedebugnameValue;
	TOKEN_PRIVILEGES tkp;
	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
	{
		return FALSE;
	}
	if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &sedebugnameValue))
	{
		CloseHandle(hToken);
		return FALSE;
	}
	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Luid = sedebugnameValue;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	if (!AdjustTokenPrivileges(hToken, FALSE, &tkp, sizeof(tkp), NULL, NULL))
	{
		CloseHandle(hToken);
		return FALSE;
	}
	return TRUE;
}

BOOL CProcess::open(DWORD pid) {
	m_hProcess = OpenProcess(PROCESS_ALL_ACCESS, false, pid);

	HANDLE hProcessSnap = NULL;
	PROCESSENTRY32 pe32 = { 0 };
	hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hProcessSnap == (HANDLE)-1)
	{
		return FALSE;
	}
	pe32.dwSize = sizeof(PROCESSENTRY32);
	if (Process32First(hProcessSnap, &pe32))
	{
		do
		{
			if (pid == pe32.th32ProcessID)
			{
				m_szProcecssName = pe32.szExeFile;
			}
		} while (Process32Next(hProcessSnap, &pe32));
	}
	else return FALSE;
	CloseHandle(hProcessSnap);
	m_nPid = pid;
	return TRUE;
}

BOOL CProcess::open(string szProcessName) 
{
	STARTUPINFO st;
	PROCESS_INFORMATION pi;
	PROCESSENTRY32 ps;
	HANDLE hSnapshot;

	ZeroMemory(&st, sizeof(STARTUPINFO));
	ZeroMemory(&pi, sizeof(PROCESS_INFORMATION));
	st.cb = sizeof(STARTUPINFO);
	ZeroMemory(&ps, sizeof(PROCESSENTRY32));
	ps.dwSize = sizeof(PROCESSENTRY32);
	hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	if (hSnapshot == INVALID_HANDLE_VALUE) return FALSE;

	if (!Process32First(hSnapshot, &ps)) return FALSE;
	do
	{
		if (lstrcmpi(ps.szExeFile, szProcessName.c_str()) == 0)
		{
			m_nPid = ps.th32ProcessID;
			m_szProcecssName = szProcessName;
			m_hProcess = OpenProcess(PROCESS_ALL_ACCESS, false, m_nPid);
			CloseHandle(hSnapshot);
		}
	} while (Process32Next(hSnapshot, &ps));
	CloseHandle(hSnapshot);
	return FALSE;
}

BOOL CProcess::open(HANDLE hProcess)
{
	DWORD pid = GetProcessId(hProcess);
	open(pid);
	CloseHandle(hProcess);
	return TRUE;
}

BOOL CProcess::terminate()
{
	if (m_hProcess == NULL) return FALSE;
	if (TerminateProcess(m_hProcess, 0)) return TRUE;
	return FALSE;
}

BOOL CProcess::suspend()
{
	if (m_hProcess == NULL) return FALSE;
	NtSuspendProcess = (_NtSuspendProcess)
		GetProcAddress(GetModuleHandle(_T("ntdll")), "NtSuspendProcess");
	EnableDebugPrivilege();
	if (NtSuspendProcess) NtSuspendProcess(m_hProcess);
	return TRUE;
}

BOOL CProcess::resume()
{
	if (m_hProcess == NULL)
	{
		return FALSE;
	}
	NtResumeProcess = (_NtResumeProcess)
		GetProcAddress(GetModuleHandle(_T("ntdll")), "NtResumeProcess");
	EnableDebugPrivilege();
	if (NtResumeProcess) NtResumeProcess(m_hProcess);
	return FALSE;
}

BOOL CProcess::inject(string szDllName)
{
	const DWORD dwThreadSize = 5 * 1024;
	DWORD dwWriteBytes;

	EnableDebugPrivilege();

	if (m_hProcess == NULL) return FALSE;

	void* pRemoteThread = VirtualAllocEx(m_hProcess,
		0,
		dwThreadSize,
		MEM_COMMIT, PAGE_EXECUTE_READWRITE);
	if (!pRemoteThread) return FALSE;

	if (!WriteProcessMemory(m_hProcess, pRemoteThread, (LPVOID)szDllName.c_str(), dwThreadSize, 0))
		return FALSE;

	LPVOID pFunc = LoadLibraryA;

	HANDLE hRemoteThread = CreateRemoteThread(m_hProcess,
		NULL,
		0,
		(LPTHREAD_START_ROUTINE)pFunc,
		pRemoteThread,
		0,
		&dwWriteBytes);
	if (!hRemoteThread) return FALSE;

	WaitForSingleObject(hRemoteThread, INFINITE);
	VirtualFreeEx(m_hProcess, pRemoteThread, dwThreadSize, MEM_COMMIT);
	CloseHandle(hRemoteThread);
	CloseHandle(m_hProcess);
	return 0;
}

BOOL CProcess::free(string szDllName)
{
	return TRUE;
}

PROCESS_MEMORY_COUNTERS CProcess::getProcessMemoryInfo()
{
	PROCESS_MEMORY_COUNTERS cmp = { 0 };
	if (m_hProcess == NULL) return cmp;
	if (GetProcessMemoryInfo(m_hProcess, &cmp, sizeof(cmp)) != 0) return cmp;
	printf("%d",GetLastError());
	printf("AWSL");
	return cmp;
}